package com.chinalcoal.ai.utils.util;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.PDPageTree;
import org.apache.pdfbox.pdmodel.PDResources;
import org.apache.pdfbox.pdmodel.graphics.image.PDImageXObject;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.pdfbox.text.TextPosition;
import org.apache.poi.openxml4j.exceptions.InvalidFormatException;
import org.apache.poi.xwpf.usermodel.*;
import org.apache.poi.util.Units;

import java.awt.image.BufferedImage;
import javax.imageio.ImageIO;
import java.io.*;
import java.util.List;

/**
 * 改进的PDF转Word工具类
 * 尝试保持内容顺序和图片
 */
public class PdfToWord {

    public static void main(String[] args) {
        String wordPath = "G:\\我的项目\\PDF转图片\\pdfToImage\\jbpm集成.docx";
        String pdfPath = "G:\\我的项目\\PDF转图片\\pdfToImage\\jbpm集成.PDF";
        try {
            PdfToWord.convertPdfToWord(
                    new File(pdfPath),
                    wordPath
            );
            System.out.println("转换完成！");
        } catch (IOException e) {
            System.err.println("转换失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 转换PDF到Word，尝试保持顺序
     *
     * @param pdfFile      输入PDF文件
     * @param outputDocxPath 输出Word文件路径
     * @throws IOException 文件操作异常
     */
    public static void convertPdfToWord(File pdfFile, String outputDocxPath) throws IOException {
        try (PDDocument pdfDocument = PDDocument.load(pdfFile);
             XWPFDocument wordDocument = new XWPFDocument()) {

            // 获取所有页面
            PDPageTree pages = pdfDocument.getPages();

            // 逐页处理
            for (int i = 0; i < pages.getCount(); i++) {
                PDPage page = pages.get(i);

                // 1. 提取文本（尝试按位置排序）
                extractTextFromPage(pdfDocument, page, i, wordDocument);

                // 2. 提取图片
                extractImagesFromPage(page, wordDocument);
            }

            // 保存Word文档
            try (FileOutputStream out = new FileOutputStream(outputDocxPath)) {
                wordDocument.write(out);
            }
        }
    }

    /**
     * 从PDF页面提取文本（尝试保持顺序）
     */
    private static void extractTextFromPage(PDDocument pdfDocument, PDPage page, int pageIndex,
                                            XWPFDocument wordDocument) throws IOException {

        // 使用普通的PDFTextStripper提取文本
        PDFTextStripper textStripper = new PDFTextStripper() {
            @Override
            protected void writeString(String text, List<TextPosition> textPositions) throws IOException {
                // 检查是否为水印文本，如果是则跳过
                if (text != null && isWatermarkText(text)) {
                    return; // 跳过水印文本
                }

                // 添加非水印文本到Word文档
                if (text != null && !text.trim().isEmpty()) {
                    XWPFParagraph paragraph = wordDocument.createParagraph();
                    XWPFRun run = paragraph.createRun();
                    run.setText(text);
                    run.setFontFamily("宋体"); // 使用中文字体
                }
            }
        };

        textStripper.setSortByPosition(true); // 关键：按位置排序
        textStripper.setStartPage(pageIndex + 1);
        textStripper.setEndPage(pageIndex + 1);
        textStripper.getText(pdfDocument);
    }

    /**
     * 检查文本是否为水印
     */
    private static boolean isWatermarkText(String text) {
        // 你需要根据实际情况修改这个判断逻辑
        // 比如，如果你知道水印文本是"机密文档 - 请勿外传"，可以这样判断：
        if ("机密文档 - 请勿外传".equals(text)) {
            return true;
        }

        // 或者包含某些关键词的文本认为是水印
        if (text.contains("机密") || text.contains("水印") || text.contains("保密") ||
            text.contains("draft") || text.contains("confidential")) {
            return true;
        }

        return false;
    }

    /**
     * 从PDF页面提取图片
     */
    private static void extractImagesFromPage(PDPage page, XWPFDocument wordDocument) throws IOException {
        PDResources resources = page.getResources();
        if (resources == null) {
            return;
        }

        // 遍历所有XObject（通常包含图片）
        Iterable<org.apache.pdfbox.cos.COSName> xObjectNames = resources.getXObjectNames();
        for (org.apache.pdfbox.cos.COSName name : xObjectNames) {
            try {
                org.apache.pdfbox.pdmodel.graphics.PDXObject xObject = resources.getXObject(name);
                if (xObject instanceof PDImageXObject) {
                    PDImageXObject image = (PDImageXObject) xObject;

                    // 将图片添加到Word
                    addImageToWordDocument(image, wordDocument);
                }
            } catch (IOException e) {
                System.err.println("处理图片时出错: " + e.getMessage());
            }
        }
    }

    /**
     * 添加图片到Word文档
     */
    private static void addImageToWordDocument(PDImageXObject image, XWPFDocument wordDocument) throws IOException {
        // 创建新段落放置图片
        XWPFParagraph paragraph = wordDocument.createParagraph();
        // 使用居中对齐，因为在PDF中图片通常居中显示
        paragraph.setAlignment(ParagraphAlignment.CENTER);
        
        // 添加一个空行，以提供更好的上下文分隔
        XWPFRun run = paragraph.createRun();
        run.addBreak(); // 添加换行
        
        // 添加图片
        // 将PDImage转换为字节数组
        BufferedImage bufferedImage = image.getImage();
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ImageIO.write(bufferedImage, "png", baos);
        byte[] imageData = baos.toByteArray();

        // 添加图片到Word
        try {
            run.addPicture(new ByteArrayInputStream(imageData),
                    XWPFDocument.PICTURE_TYPE_PNG,
                    "image.png",
                    Units.toEMU(image.getWidth()),
                    Units.toEMU(image.getHeight()));
            // 在图片后添加换行，确保后续内容不会紧贴图片
            run.addBreak();
        } catch (InvalidFormatException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 处理加密的PDF
     */
    public static void convertEncryptedPdf(File pdfFile, String password, String outputDocxPath) throws IOException {
        try (PDDocument pdfDocument = PDDocument.load(pdfFile, password)) {
            // 其余逻辑与普通转换相同
            convertPdfToWord(pdfFile, outputDocxPath);
        }
    }
}
